import datetime
import tushare as ts

ts.set_token("9178f4653b48b9f2b86e50aedec71960f76540cd96effe1ed9360dea")
pro = ts.pro_api()


def getDayStock(stockCode: str, startDate: str, endDate: str):
    """
    获取股票或指数的日线数据，根据 stockCode 自动选择接口。
    :param stockCode: 股票或指数代码
    :param startDate: 开始日期 (格式：YYYYMMDD)
    :param endDate: 结束日期 (格式：YYYYMMDD)
    :return: 对应的日线数据
    """
    # 判断接口类型
    if stockCode.endswith(".SH") or stockCode.endswith(".SZ"):
        if stockCode.startswith("000") or stockCode.startswith("399"):
            # 指数代码，调用指数日线接口
            return pro.index_daily(
                ts_code=stockCode, start_date=startDate, end_date=endDate
            )
        else:
            # 股票代码，调用股票日线接口
            return pro.daily(ts_code=stockCode, start_date=startDate, end_date=endDate)
    elif (
        stockCode.endswith(".BJ")
        or stockCode.endswith(".HK")
        or stockCode.endswith(".US")
    ):
        # 北京证券交易所、港股、美股直接调用股票接口
        return pro.daily(ts_code=stockCode, start_date=startDate, end_date=endDate)
    else:
        raise ValueError(f"无法识别的股票代码格式：{stockCode}")


# 使用示例
if __name__ == "__main__":
    # 查询股票交易明细并且添加
    today = datetime.today()
    endTime = today.strftime("%Y%m%d")
    startTime = (today - datetime.timedelta(days=180)).strftime("%Y%m%d")
    daily_data = getDayStock("600519.SH", startTime, endTime)
    AntLogger.infodaily_data)
